package com.hyc.rateLimit;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author 冷环渊
 * 用原子请求来限制异步
 * @date 2025/6/13 13:22
 * @description AsyncTokenBucketRateLimiter
 */
public class AsyncTokenBucketRateLimiter implements RateLimit {
    // 令牌生成速率 (毫秒/个)
    private final long requestPerSecond;
    // 桶容量
    private final int maxToken;

    // 当前令牌数 (原子类型)
    private AtomicInteger tokens;
    // 最后填充时间 (原子类型)
    private AtomicLong lastRefillTime;

    public AsyncTokenBucketRateLimiter(int requestPerSecond, int maxToken) {
        this.requestPerSecond = 1000L * requestPerSecond;
        this.maxToken = maxToken;
        this.tokens = new AtomicInteger(maxToken);
        this.lastRefillTime = new AtomicLong(System.currentTimeMillis());
    }

    //这个方法用于获取是否能够执行请求 并且是否能够获取令牌
    public synchronized boolean getToken(int n) {
        //当前时间
        long now = System.currentTimeMillis();
        //获取到当前时间的间隔并且替换当前时间
        long timePassed = now - lastRefillTime.getAndSet(now);
        //计算出需要增加的 token
        int tokensToAdd = (int) (timePassed / requestPerSecond);
        //用当前时间能生成的token和桶中还能放做少token取小 来以防溢出
        tokens.addAndGet(Math.min(tokensToAdd, maxToken - tokens.get()));
        return tokens.getAndAdd(-n) >= n;
    }


    @Override
    public boolean getToken() {
        return getToken(1);
    }
}
